package com.grf.sse.utils;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

/**
 * @description:
 * @author: guorf
 * @date: 2024/8/7
 */
@Slf4j
public class SseEmitterUtil {
    /**
     * 当前连接数
     */
    private static final AtomicInteger count = new AtomicInteger(0);

    /**
     * 使用map对象，便于根据userId来获取对应的SseEmitter，或者放redis里面
     */
    private static final Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    /**
     * @description: 建立连接
     * @author: guorf
     * @date: 2024/8/7
     */
    public static SseEmitter connect(String userId) {
        // 设置超时时间，0表示不过期。默认30秒，超过时间未完成会抛出异常：AsyncRequestTimeoutException
        SseEmitter sseEmitter = new SseEmitter(0L);

        // 注册回调
        sseEmitter.onCompletion(completionCallBack(userId));
        sseEmitter.onError(errorCallBack(userId));
        sseEmitter.onTimeout(timeoutCallBack(userId));
        sseEmitterMap.put(userId, sseEmitter);
        // 数量+1
        count.getAndIncrement();
        log.info("创建新的sse连接，当前用户：{}", userId);
        return sseEmitter;
    }
    /**
     * @description: 发送消息给指定的人
     * @author: guorf
     * @date: 2024/8/7
     */
    public static void sendMessage(String userId, String message) {
        try {
            sseEmitterMap.get(userId).send(message);
        } catch (Exception e) {
            log.error("用户【{}】消息发送异常：{}", userId, e.getMessage());
            // 移除用户
            removeUser(userId);
        }
    }
    /**
     * @description: 群发指定用户
     * @author: guorf
     * @date: 2024/8/7
     */
    public static void batchSendMessage(List<String> ids, String message){
        ids.forEach(userId -> sendMessage(userId, message));
    }

    /**
     * @description: 群发所有人
     * @author: guorf
     * @date: 2024/8/7
     */
    public static void batchSendMessage(String message) {
        sseEmitterMap.forEach((k, v) -> {
            try {
                v.send(message, MediaType.APPLICATION_JSON);
            } catch (Exception e) {
                log.error("用户【{}】消息发送异常：{}", k, e.getMessage());
                // 移除用户
                removeUser(k);
            }
        });
    }


    /**
     * 获取当前连接信息
     */
    public static List<String> getIds() {
        return new ArrayList<>(sseEmitterMap.keySet());
    }

    /**
     * @description: 移除指定用户
     * @author: guorf
     * @date: 2024/8/7
     */
    public static void removeUser(String userId) {
        sseEmitterMap.remove(userId);
        // 数量-1
        count.getAndDecrement();
    }

    /**
     * @description: 获取在线用户数
     * @author: guorf
     * @date: 2024/8/7
     */
    public static int getUserCount() {
        return count.intValue();
    }

    /**
     * @description: 超时回调
     * @author: guorf
     * @date: 2024/8/7
     */
    private static Runnable timeoutCallBack(String userId) {
        return () -> {
            log.info("用户：{}连接超时", userId);
            removeUser(userId);
        };
    }
    /**
     * @description: 连接错误回调
     * @author: guorf
     * @date: 2024/8/7
     */
    private static Consumer<Throwable> errorCallBack(String userId) {
        return throwable -> {
            log.info("用户：{}连接异常", userId);
            removeUser(userId);
        };
    }
    /**
     * @description: 结束连接回调
     * @author: guorf
     * @date: 2024/8/7
     */
    private static Runnable completionCallBack(String userId) {
        return () -> {
            log.info("用户：{}结束连接", userId);
            removeUser(userId);
        };
    }
}
